#pragma once
#include <iostream>
#include <utility>
using namespace std;

namespace old
{

enum COLOR
{
    RED,
    BLACK,
};

template<class K, class V>
struct rbt_node
{
    rbt_node<K, V>* _left;
    rbt_node<K, V>* _right;
    rbt_node<K, V>* _parent;

    pair<K, V> _kv;
    COLOR _col;

    rbt_node<K, V>(const pair<K, V>& kv)
        : _left(nullptr), _right(nullptr), _parent(nullptr)
        , _kv(kv), _col(RED)
    {}
};

template<class K, class V>
class rb_tree
{
public:
    typedef rbt_node<K, V> node;

public:
    rb_tree<K, V>()
    {}

    rb_tree<K, V>(const rb_tree<K, V>& t)
    {
        _root = copy(_root);
    }

    rb_tree<K, V>& operator=(const rb_tree<K, V> t)
    {
        if (this != &t)
        {
            swap(_root, t._root);
        }
        return *this;
    }

    ~rb_tree<K, V>()
    {
        destroy(_root);
        _root = nullptr;
    }

private:
    node* copy(node* root)
    {
        if (!root)
            return nullptr;

        node* new_node = new node(root->_kv);
        new_node->_left = copy(root->_left);
        new_node->_right = copy(root->_right);
        return new_node;
    }

    void destroy(node* root)
    {
        if (!root)
            return;

        destroy(root->_left);
        destroy(root->_right);
        delete root;
    }

public:
    node* find(const K& key)
    {
        node* cur = _root;

        while (cur)
        {
            if (cur->_kv.first < key)
                cur = cur->_right;
            else if (cur->_kv.first > key)
                cur = cur->_left;
            else
                return cur;
        }

        return nullptr;
    }

    bool insert(const pair<K, V>& kv)
    {
        if (!_root)
        {
            _root = new node(kv);
            _root->_col = BLACK;
            return true;
        }

        node* parent = nullptr;
        node* cur = _root;

        while (cur)
        {
            if (cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        }

        cur = new node(kv);
        if (parent->_kv.first < kv.first)
            parent->_right = cur;
        else
            parent->_left = cur;
        cur->_parent = parent;

        while (parent && parent->_col == RED)
        {
            node* grandpa = parent->_parent;

            if (grandpa->_left == parent)
            {
                node* uncle = grandpa->_right;

                if (uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandpa->_col = RED;

                    cur = grandpa;
                    parent = cur->_parent;
                }
                else
                {
                    if (parent->_left == cur)
                    {
                        rotate_r(grandpa);
                        cur->_col = grandpa->_col = RED;
                        parent->_col = BLACK;
                    }
                    else
                    {
                        rotate_l(parent);
                        rotate_r(grandpa);
                        grandpa->_col = parent->_col = RED;
                        cur->_col = BLACK;
                    }
                    break;
                }
            }
            else
            {
                node* uncle = grandpa->_left;

                if (uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandpa->_col = RED;

                    cur = grandpa;
                    parent = cur->_parent;
                }
                else
                {
                    if (parent->_left == cur)
                    {
                        rotate_r(parent);
                        rotate_l(grandpa);
                        grandpa->_col = parent->_col = RED;
                        cur->_col = BLACK;
                    }
                    else
                    {
                        rotate_l(grandpa);
                        cur->_col = grandpa->_col = RED;
                        parent->_col = BLACK;
                    }
                    break;
                }
            }
        }

        _root->_col = BLACK;
        return true;
    }

private:
    void rotate_l(node* parent)
    {
        node* subr = parent->_right;
        node* subrl = subr->_left;

        parent->_right = subrl;
        if (subrl) subrl->_parent = parent;

        node* pparent = parent->_parent;

        subr->_left = parent;
        parent->_parent = subr;

        if (parent == _root)
            _root = subr;
        else
        {
            if (pparent->_left == parent)
                pparent->_left = subr;
            else
                pparent->_right = subr;
        }
        subr->_parent = pparent;
    }

    void rotate_r(node* parent)
    {
        node* subl = parent->_left;
        node* sublr = subl->_right;

        parent->_left = sublr;
        if (sublr) sublr->_parent = parent;

        node* pparent = parent->_parent;

        subl->_right = parent;
        parent->_parent = subl;

        if (parent == _root)
        {
            _root = subl;
            subl->_parent = nullptr;
        }
        else
        {
            if (pparent->_left == parent)
                pparent->_left = subl;
            else
                pparent->_right = subl;
            subl->_parent = pparent;
        }
    }

public:
    void inorder()
    {
        _inorder(_root);
        cout << endl;
    }

    bool is_rbtree()
    {
        if (_root && _root->_col == RED)
            return false;

        int mark = -1;
        return check(_root, mark, 0);
    }

private:
    void _inorder(node* root)
    {
        if (!root) return;

        _inorder(root->_left);
        cout << root->_kv.first << ":" << root->_kv.second << " ";
        _inorder(root->_right);
    }

    bool check(node* root, int& mark, int cnt)
    {
        if (!root)
        {
            if (mark == -1)
                mark = cnt;
            else if (mark != cnt)
            {
                cout << "block nodes count error\n"
                     << endl;
                return false;
            }

            return true;
        }

        if (root->_col == BLACK)
            cnt++;

        if (root->_col == RED && root->_parent->_col == RED)
        {
            cout << "consecutive red nodes "
                 << root->_parent->_kv.first << " and "
                 << root->_kv.first << endl;
            return false;
        }

        return check(root->_left, mark, cnt)
               && check(root->_right, mark, cnt);
    }

private:
    node* _root = nullptr;
};

}